#########################################
# Author: FuHuizhong<fuhuizn@gmail.com> #
# Licence: GPL3                         #
# Port Mapping Tool with UPNP           #
#########################################
import socket
import thread
import re
import urllib2
import xml.dom.minidom
import time,sys
import base64

s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.bind( ('',0) )

ssdp_req = 'M-SEARCH * HTTP/1.1\r\nHOST: 239.255.255.250:1900\r\nMAN:"ssdp:discover"\r\nMX:3\r\nST:UPnP:rootdevice\r\n\r\n'
s.sendto( ssdp_req,('239.255.255.250',1900) )

#auth_enc = 'YWRtaW46YWRtaW4='
def newRequest( url ):
    req = urllib2.Request( url,unverifiable=True )
    info = urllib2.urlparse.urlparse( url )
    req.add_header( 'HOST', info[1] )
    #global auth_enc
    #req.add_header( 'Authorization', 'Basic ' + auth_enc )
    return req
    
def is_ok(data):
    try:
        s1,s2 = data.split('\r\n',1)
    except:
        s1 = ''
    if s1.endswith('200 OK'):
        return True
    else:
        return False

location = ''
def get_location(data):
    global location
    #print data
    ret = re.findall( '[Ll][Oo][Cc][Aa][Tt][Ii][Oo][Nn]: ([^\r]*)',data )
    location = ret[0]
    print 'discription:',location

ctrl_url = ''
serviceType = ''
def get_ctrlurl():
    global location,ctrl_url,serviceType
    fp = urllib2.urlopen( newRequest( location ) )
    #print fp.headers
    data = ''
    for buf in fp:
        data += buf
    fp.close()
    #print data
    igd = xml.dom.minidom.parseString(data)
    #get control url from the xml
    for service in igd.getElementsByTagName('service'):
        val = {}
        for child in service.childNodes:
            try:
                val[child.nodeName.lower()] = child.childNodes[0].nodeValue
            except:
                pass
        if (val['servicetype'].lower()=='urn:schemas-upnp-org:service:wanpppconnection:1') or (val['servicetype'].lower()=='urn:schemas-upnp-org:service:wanipconnection:1'):
            ctrl_url = urllib2.urlparse.urljoin( location,val['controlurl'] )
            serviceType = val['servicetype']
            break
    ret = get_port_map( 1 )
    print 'controlurl:',ctrl_url

def get_local_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(('google.com', 0))
    ip,port = s.getsockname()
    return ip
 
def add_port(ExternalPort,Protocol,InternalPort):
    global ctrl_url,serviceType
    if ctrl_url == '':
        return ''
    InternalClient = get_local_ip()
    actionName = 'AddPortMapping'
    actionParams = '<NewRemoteHost></NewRemoteHost>\r\n\
<NewExternalPort>%d</NewExternalPort>\r\n\
<NewProtocol>%s</NewProtocol>\r\n\
<NewInternalPort>%d</NewInternalPort>\r\n\
<NewInternalClient>%s</NewInternalClient>\r\n\
<NewEnabled>1</NewEnabled>\r\n\
<NewPortMappingDescription>upnp.py</NewPortMappingDescription>\r\n\
<NewLeaseDuration>0</NewLeaseDuration>\r\n' % (ExternalPort,Protocol,InternalPort,InternalClient)
    control_doc = '<?xml version="1.0" encoding="utf-8"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\r\n\
<s:Body>\r\n\
<u:%s xmlns:u="%s">\r\n\
%s</u:%s>\r\n\
</s:Body>\r\n\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
    up = urllib2.urlparse.urlparse(ctrl_url)
    header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
    #global auth_enc
    #header += 'Authorization: Basic ' + auth_enc + '\r\n'
    header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
    header += 'CONTENT-TYPE: text/xml\r\n'
    header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n\r\n'
    req = header + control_doc
    host,port = up[1].split(':')
    fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    fd.connect( (host,int(port)) )
    fd.send( req )
    data = ''
    while 1:
        buf = fd.recv(1492)
        if len(buf)>0:
            data += buf
        else:
            break
    fd.close()
    if is_ok(data):
        print 'OK ADD',InternalClient,Protocol,str(ExternalPort)
    else:
        print 'Fail'
    return data

def del_port(ExternalPort,Protocol):
    global ctrl_url,serviceType
    if ctrl_url == '':
        return ''
    actionName = 'DeletePortMapping'
    actionParams = '<NewRemoteHost></NewRemoteHost>\r\n\
<NewExternalPort>%d</NewExternalPort>\r\n\
<NewProtocol>%s</NewProtocol>\r\n' % (ExternalPort,Protocol)
    control_doc = '<?xml version="1.0" encoding="utf-8"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\r\n\
<s:Body>\r\n\
<u:%s xmlns:u="%s">\r\n\
%s</u:%s>\r\n\
</s:Body>\r\n\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
    up = urllib2.urlparse.urlparse(ctrl_url)
    header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
    #global auth_enc
    #header += 'Authorization: Basic ' + auth_enc + '\r\n'
    header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
    header += 'CONTENT-TYPE: text/xml\r\n'
    header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n\r\n'
    req = header + control_doc
    host,port = up[1].split(':')
    fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    fd.connect( (host,int(port)) )
    fd.send( req )
    data = ''
    while 1:
        buf = fd.recv(1492)
        if len(buf)>0:
            data += buf
        else:
            break
    fd.close()
    if is_ok(data):
        print 'OK'
    else:
        print 'Fail'
    return data

def get_port_map(PortMappingIndex):
    global ctrl_url,serviceType
    if ctrl_url == '':
        return ''
    actionName = 'GetGenericPortMappingEntry'
    actionParams = '<NewPortMappingIndex>%d</NewPortMappingIndex>\r\n\
<NewRemoteHost></NewRemoteHost>\r\n\
<NewExternalPort></NewExternalPort>\r\n\
<NewProtocol></NewProtocol>\r\n\
<NewInternalPort></NewInternalPort>\r\n\
<NewInternalClient></NewInternalClient>\r\n\
<NewEnabled>1</NewEnabled>\r\n\
<NewPortMappingDescription></NewPortMappingDescription>\r\n\
<NewLeaseDuration></NewLeaseDuration>\r\n' % (PortMappingIndex)
    control_doc = '<?xml version="1.0" encoding="utf-8"?>\r\n\
<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/" s:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\r\n\
<s:Body>\r\n\
<u:%s xmlns:u="%s">\r\n\
%s</u:%s>\r\n\
</s:Body>\r\n\
</s:Envelope>\r\n' % (actionName,serviceType,actionParams,actionName)
    up = urllib2.urlparse.urlparse(ctrl_url)
    header = 'POST %s HTTP/1.1\r\nHost: %s\r\n' %( up[2],up[1] )
    #global auth_enc
    #header += 'Authorization: Basic ' + auth_enc + '\r\n'
    header += 'SOAPACTION: "' + serviceType + '#'+ actionName +'"\r\n'
    header += 'CONTENT-TYPE: text/xml\r\n'
    header += 'Content-Length: '+ str( len( control_doc ) ) + '\r\n\r\n'
    req = header + control_doc
    host,port = up[1].split(':')
    fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    fd.connect( (host,int(port)) )
    fd.send( req )
    data = ''
    while 1:
        buf = fd.recv(1492)
        if len(buf)>0:
            data += buf
        else:
            break
    fd.close()
    return data
    
def output_info(data,i):
    port = re.findall('<NewExternalPort>([0-9].*)</NewExternalPort>',data)[0]
    protocol = re.findall('<NewProtocol>([a-zA-Z].*)</NewProtocol>',data)[0]
    client = re.findall('<NewInternalClient>([^<>].*)</NewInternalClient>',data)[0]
    print i,client,protocol,port

def recv(s):
    while 1:
        try:
            data,addr = s.recvfrom(4096)
            get_location(data)
            get_ctrlurl()
        except:
            pass
        
id = thread.start_new_thread( recv,(s,) )

def action( c ):
    if c.startswith('q'):
        s.close()
        sys.exit()
    try:
        if c.startswith('r'):
            s.sendto( ssdp_req,('239.255.255.250',1900) )
            return 0
        if c.startswith('a'):
            argv = c.split()
            if argv[1]=='s':
                port = int(argv[3])
                ret = add_port(port,argv[2],port)
                return 0
            if argv[1]=='r':
                for port in range( int(argv[3]), int(argv[4])+1 ):
                    ret = add_port(port,argv[2],port)
                return 0
            if argv[1]=='m':
                for i in range( 3,len(argv) ):
                    port = int( argv[i] )
                    ret = add_port(port,argv[2],port)
                return 0
        if c.startswith('d'):
            argv = c.split()
            if argv[1]=='s':
                port = int(argv[3])
                ret = del_port(port,argv[2])
                return 0
            if argv[1]=='r':
                for port in range( int(argv[3]), int(argv[4])+1 ):
                    ret = del_port(port,argv[2])
                return 0
            if argv[1]=='m':
                for i in range( 3,len(argv) ):
                    port = int( argv[i] )
                    ret = del_port(port,argv[2])
                return 0
        if c.startswith('p'):
            if c.strip()=='p':
                i = 0
                while 1:
                    try:
                        ret = get_port_map(i)
                        i += 1
                        if is_ok(ret):
                            output_info(ret,i-1)
                        else:
                            break
                    except:
                        pass
            else:
                pass
            try:
                p,index = c.split()
                index = int(index)
                ret = get_port_map(index)
                if is_ok(ret):
                    output_info(ret,index)
                else:
                    print 'No info'
            except:
                pass
            return 0
    except Exception,msg:
        print msg
        return 0

time.sleep(1)
try:    
    rc = open('upnprc','r')
    while 1:
        buf = rc.readline()
        if len(buf)<=0:
            break
        buf = buf.strip()
        action( buf )
except:
    pass
while 1:
    c = raw_input('>')
    action( c )
    time.sleep(1)
    
